package com.example.homeworktagger.service;

import com.example.homeworktagger.model.Question;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.EmbeddingModel;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.MutationResult;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.collection.HasCollectionParam;
import io.milvus.param.collection.DropCollectionParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.index.CreateIndexParam;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;

/**
 * 数据导入服务
 * 
 * 负责将题目数据导入到Milvus向量数据库中，包括：
 * - 创建集合
 * - 构建题目向量
 * - 批量导入数据
 * - 创建索引
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class DataImportService {

    /** Milvus客户端 */
    private final MilvusServiceClient milvusClient;
    
    /** 向量嵌入模型 */
    private final EmbeddingModel embeddingModel;

    /** 集合名称 */
    private static final String COLLECTION_NAME = "questions";
    
    /** 向量字段名称 */
    private static final String VECTOR_FIELD = "question_vector";
    
    /** 问题文本字段名称 */
    private static final String QUESTION_FIELD = "question";
    
    /** 答案文本字段名称 */
    private static final String ANSWER_FIELD = "answer";
    
    /** 向量维度 */
    private static final int VECTOR_DIMENSION = 384;

    /**
     * 导入题目列表到Milvus
     * 
     * @param questions 题目列表
     */
    public void importQuestions(List<Question> questions) {
        try {
            // 1. 创建或重建集合
            createCollection();

            // 2. 准备批量导入数据
            List<InsertParam.Field> fields = prepareFields(questions);

            // 3. 执行批量插入
            executeInsert(fields);

            // 4. 创建索引
            createIndex();
            
            log.info("成功导入 {} 条问题", questions.size());
        } catch (Exception e) {
            log.error("导入数据失败", e);
            throw new RuntimeException("导入数据失败: " + e.getMessage(), e);
        }
    }

    /**
     * 准备批量导入的字段数据
     * 
     * @param questions 题目列表
     * @return 准备好的字段数据
     */
    private List<InsertParam.Field> prepareFields(List<Question> questions) {
        List<InsertParam.Field> fields = new ArrayList<>();
        List<String> questionList = new ArrayList<>();
        List<String> answerList = new ArrayList<>();
        List<List<Float>> vectorList = new ArrayList<>();

        // 为每个题目生成向量表示
        for (Question question : questions) {
            // 生成向量
            Embedding embedding = embeddingModel.embed(question.getQuestion()).content();
            List<Float> vector = new ArrayList<>();
            for (float f : embedding.vector()) {
                vector.add(f);
            }

            questionList.add(question.getQuestion());
            answerList.add(question.getAnswer());
            vectorList.add(vector);
        }

        // 添加字段
        fields.add(new InsertParam.Field(QUESTION_FIELD, questionList));
        fields.add(new InsertParam.Field(ANSWER_FIELD, answerList));
        fields.add(new InsertParam.Field(VECTOR_FIELD, vectorList));
        
        return fields;
    }

    /**
     * 执行批量插入操作
     * 
     * @param fields 字段数据
     */
    private void executeInsert(List<InsertParam.Field> fields) {
        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .withFields(fields)
                .build();

        R<MutationResult> insertResult = milvusClient.insert(insertParam);
        if (insertResult.getStatus() != R.Status.Success.getCode()) {
            throw new RuntimeException("数据插入失败: " + insertResult.getMessage());
        }
        
        log.info("数据插入成功");
    }

    /**
     * 创建Milvus集合
     * 首先检查集合是否存在，如果存在则先删除
     */
    public void createCollection() {
        // 检查集合是否存在
        R<Boolean> hasCollection = milvusClient.hasCollection(HasCollectionParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .build());
        
        if (hasCollection.getData()) {
            milvusClient.dropCollection(DropCollectionParam.newBuilder()
                    .withCollectionName(COLLECTION_NAME)
                    .build());
            log.info("已删除现有集合: {}", COLLECTION_NAME);
        }

        // 创建集合
        CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .withFieldTypes(List.of(
                        FieldType.newBuilder()
                                .withName("id")
                                .withDescription("主键")
                                .withDataType(io.milvus.grpc.DataType.Int64)
                                .withAutoID(true)
                                .withPrimaryKey(true)
                                .build(),
                        FieldType.newBuilder()
                                .withName(QUESTION_FIELD)
                                .withDescription("问题文本字段")
                                .withDataType(io.milvus.grpc.DataType.VarChar)
                                .withMaxLength(65535)
                                .build(),
                        FieldType.newBuilder()
                                .withName(ANSWER_FIELD)
                                .withDescription("答案文本字段")
                                .withDataType(io.milvus.grpc.DataType.VarChar)
                                .withMaxLength(65535)
                                .build(),
                        FieldType.newBuilder()
                                .withName(VECTOR_FIELD)
                                .withDescription("问题的向量表示")
                                .withDataType(io.milvus.grpc.DataType.FloatVector)
                                .withDimension(VECTOR_DIMENSION)
                                .build()
                ))
                .build();

        R<RpcStatus> createResult = milvusClient.createCollection(createCollectionParam);
        if (createResult.getStatus() != R.Status.Success.getCode()) {
            throw new RuntimeException("创建集合失败: " + createResult.getMessage());
        }
        
        log.info("集合创建成功: {}", COLLECTION_NAME);
    }

    /**
     * 为向量字段创建索引
     */
    private void createIndex() {
        CreateIndexParam createIndexParam = CreateIndexParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .withFieldName(VECTOR_FIELD)
                .withIndexType(IndexType.IVF_FLAT)
                .withMetricType(MetricType.COSINE)
                .withExtraParam("{\"nlist\":1024}")
                .build();

        R<RpcStatus> createIndexResult = milvusClient.createIndex(createIndexParam);
        if (createIndexResult.getStatus() != R.Status.Success.getCode()) {
            throw new RuntimeException("创建索引失败: " + createIndexResult.getMessage());
        }
        
        log.info("索引创建成功");
    }
} 